#pragma once

#include <mutex>
#include <memory>
#include <condition_variable>

// 基于单链表实现的细粒度锁的并发队列
template<typename T>
class threadsafe_que
{
private:
    struct node
    {
        std::shared_ptr<T> data;
        std::unique_ptr<node> next;
    };

    std::unique_ptr<node> head;
    std::mutex head_mutex;
    node* tail;
    std::mutex tail_mutex;
    std::condition_variable data_cond;

    node* get_tail()
    {
        std::lock_guard<std::mutex> lk(tail_mutex);
        return tail;
    }

    // 需要在持有 head_mutex 时调用
    std::unique_ptr<node> pop_head()
    {
        std::unique_ptr<node> old_head = std::move(head);
        head = std::move(old_head->next);
        return old_head;
    }

    // 等待队列有数据并转移锁的使用权
    std::unique_lock<std::mutex> wait_for_data()
    {
        std::unique_lock<std::mutex> head_lock(head_mutex);
        data_cond.wait(head_lock, [this]{ return head.get() != this->get_tail(); });
        return head_lock;
    }

    std::unique_ptr<node> wait_pop_head()
    {
        std::unique_lock<std::mutex> head_lock(wait_for_data());
        return pop_head();
    }

    void wait_pop_head(T& value)
    {
        std::unique_lock<std::mutex> head_lock(wait_for_data());
        value = std::move(*(head->data));
        pop_head();
    }


public:
    threadsafe_que() : head(new node), tail(head.get()) {}
    threadsafe_que(const threadsafe_que&) = delete;
    threadsafe_que& operator=(const threadsafe_que&) = delete;

    std::shared_ptr<T> try_pop()
    {
        // 注意这里必须要先获得 head 的锁, 以免被其它线程更改
        std::lock_guard<std::mutex> head_lock(head_mutex);
        if (head.get() == get_tail())
        {
            return std::shared_ptr<T>();
        }

        std::unique_ptr<node> old_head = pop_head();
        return old_head->data;
    }

    void push(T value)
    {
        std::shared_ptr<T> new_data(std::make_shared<T>(std::move(value)));
        std::unique_ptr<node> p(new node);
        
        // 初始化节点可以不用加锁
        {
            std::lock_guard<std::mutex> tail_lock(tail_mutex);
            tail->data = new_data;
            node* new_tail = p.get();
            tail->next = std::move(p);
            tail = new_tail;
        }
        data_cond.notify_one();
    }

    std::shared_ptr<T> wait_and_pop()
    {
        std::unique_ptr<node> old_head = wait_pop_head();
        return old_head->data;
    }

    void wait_and_pop(T& value)
    {
        wait_pop_head(value);
    }

    bool empty()
    {
        std::lock_guard<std::mutex> head_lock(head_mutex);
        return head.get() == get_tail();
    }
};